#!/usr/bin/env python3
# ui/map_widget.py
# 地图小部件：渲染地图、节点、机器人、信念
from PyQt5.QtWidgets import QWidget
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QBrush, QFont, QColor
from PyQt5.QtCore import Qt, QRect
import numpy as np
from math import dist

import logging

logger = logging.getLogger(__name__)


class GraphMapWidget(QWidget):
    def __init__(self, map_manager=None, parent=None):
        super().__init__(parent)
        self.map_manager = map_manager
        self.robots = []
        self.belief_pr = None
        self.belief_pos = None
        self.scale_factor = 1.0
        self.offset_x = 0
        self.offset_y = 0
        self.capture_count = 0
        # 地图相关
        self.map_pixmap = None
        self.map_width = 0
        self.map_height = 0
        self.resolution = 0.05
        self.origin = [0.0, 0.0]
        self.map_offset_x = 0
        self.map_offset_y = 0
        self.map_scale_factor = 1.0
        # 角色颜色
        self.role_colors = {
            "evader": QColor(255, 0, 0),
            "pursuer": QColor(0, 0, 255),
        }
        self.setMinimumSize(800, 600)
        self.setMouseTracking(True)
        self.load_map()
        if (
            self.map_pixmap is None
            and self.map_manager
            and self.map_manager.node_coordinates
        ):
            self.update_graph_view()
        else:
            self.update_map_view()

    def load_map(self):
        """加载地图"""
        if self.map_manager is None:
            return
        self.map_manager.qpix_generate()
        if self.map_manager.map_pixmap is None:
            return
        self.map_pixmap = self.map_manager.get_obstacle_pixmap()
        self.map_width = self.map_manager.map_width
        self.map_height = self.map_manager.map_height
        self.resolution = self.map_manager.resolution
        self.origin = self.map_manager.origin
        logger.info("地图加载到widget")

    def reset(self):
        """重置"""
        self.load_map()
        if (
            self.map_pixmap is None
            and self.map_manager
            and self.map_manager.node_coordinates
        ):
            self.update_graph_view()
        else:
            self.update_map_view()
        self.update()

    def update_map_view(self):
        """更新地图视图"""
        if self.map_pixmap is None:
            return
        map_w = self.map_pixmap.width()
        map_h = self.map_pixmap.height()
        scale_x = self.width() / map_w
        scale_y = self.height() / map_h
        self.map_scale_factor = min(scale_x, scale_y)
        scaled_w = map_w * self.map_scale_factor
        scaled_h = map_h * self.map_scale_factor
        self.map_offset_x = (self.width() - scaled_w) / 2
        self.map_offset_y = (self.height() - scaled_h) / 2

    def update_graph_view(self):
        """更新图视图"""
        if not self.map_manager or not self.map_manager.node_coordinates:
            return
        coords = self.map_manager.node_coordinates
        min_x = min(x for x, y in coords)
        max_x = max(x for x, y in coords)
        min_y = min(y for x, y in coords)
        max_y = max(y for x, y in coords)
        width_ratio = (self.width() - 100) / (max_x - min_x) if max_x - min_x > 0 else 1
        height_ratio = (
            (self.height() - 100) / (max_y - min_y) if max_y - min_y > 0 else 1
        )
        self.scale_factor = min(width_ratio, height_ratio)
        self.offset_x = (
            self.width() - (max_x - min_x) * self.scale_factor
        ) / 2 - min_x * self.scale_factor
        self.offset_y = (
            self.height() - (max_y - min_y) * self.scale_factor
        ) / 2 - min_y * self.scale_factor

    def world_to_pixel(self, x, y):
        """世界坐标到像素"""
        if self.map_pixmap is not None and self.map_manager:
            px, py = self.map_manager.world_to_pixel(x, y)
            wx = self.map_offset_x + px * self.map_scale_factor
            wy = self.map_offset_y + py * self.map_scale_factor
            return int(wx), int(wy)
        else:
            if not self.map_manager or not self.map_manager.node_coordinates:
                return int(x), int(y)
            scaled_x = x * self.scale_factor + self.offset_x
            scaled_y = y * self.scale_factor + self.offset_y
            flipped_y = self.height() - scaled_y
            return int(scaled_x), int(flipped_y)

    def paintEvent(self, event):
        """绘制事件"""
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        # 绘制底图
        if self.map_pixmap is not None:
            self.update_map_view()
            map_w = self.map_pixmap.width()
            map_h = self.map_pixmap.height()
            scaled_w = map_w * self.map_scale_factor
            scaled_h = map_h * self.map_scale_factor
            target_rect = QRect(
                int(self.map_offset_x),
                int(self.map_offset_y),
                int(scaled_w),
                int(scaled_h),
            )
            painter.drawPixmap(target_rect, self.map_pixmap)
        else:
            painter.fillRect(self.rect(), QColor(240, 240, 240))
            if not self.map_manager or not self.map_manager.node_coordinates:
                painter.setPen(QPen(Qt.black))
                painter.setFont(QFont("Arial", 12))
                painter.drawText(self.rect(), Qt.AlignCenter, "没有地图数据可用")
                return
        # 绘制图和机器人
        if self.map_manager:
            node_num = self.map_manager.node_num
            if node_num > 0:
                node_radius = 4
                edge_width = 1
                for i in range(node_num):
                    x, y = self.map_manager.node_coordinates[i]
                    px, py = self.world_to_pixel(x, y)
                    painter.setPen(QPen(QColor(80, 80, 80), 1))
                    painter.setBrush(QBrush(Qt.NoBrush))
                    painter.drawEllipse(
                        px - node_radius,
                        py - node_radius,
                        node_radius * 2,
                        node_radius * 2,
                    )
                    for j in range(i + 1, node_num):
                        if self.map_manager.adj_matrix[i][j] == 1:
                            jx, jy = self.map_manager.node_coordinates[j]
                            pjx, pjy = self.world_to_pixel(jx, jy)
                            painter.setPen(QPen(QColor(180, 180, 180), edge_width))
                            painter.drawLine(px, py, pjx, pjy)
                # 信念
                if (
                    self.belief_pr is not None
                    and self.belief_pos is not None
                    and len(self.belief_pr) == node_num
                ):
                    belief_base_radius = 4
                    belief_scale = 8
                    for i in range(node_num):
                        if self.belief_pos[i]:
                            x, y = self.map_manager.node_coordinates[i]
                            px, py = self.world_to_pixel(x, y)
                            pr = self.belief_pr[i]
                            gray_intensity = int(200 - pr * 200)
                            belief_color = QColor(gray_intensity, 255, gray_intensity)
                            radius = belief_base_radius + int(pr * belief_scale)
                            painter.setPen(QPen(belief_color.darker(), 1))
                            painter.setBrush(QBrush(belief_color))
                            painter.drawEllipse(
                                px - radius, py - radius, radius * 2, radius * 2
                            )
                            if pr >= 0.1:
                                painter.setPen(QPen(Qt.black))
                                painter.setFont(QFont("Arial", 6))
                                painter.drawText(px + radius + 2, py + 2, f"{pr:.1f}")
                # 机器人
                for robot in self.robots:
                    if (
                        "position" not in robot
                        or robot["position"] is None
                        or "node" not in robot
                        or robot["node"] is None
                    ):
                        continue
                    try:
                        node_id = int(robot["node"])
                        if 0 <= node_id < node_num:
                            x, y = self.map_manager.node_coordinates[node_id]
                        else:
                            pos_str = robot["position"].strip("()")
                            x, y = map(float, pos_str.split(","))
                        px, py = self.world_to_pixel(x, y)
                    except:
                        continue
                    role = robot.get("role", "unknown")
                    color = self.role_colors.get(role, QColor(128, 128, 128))
                    robot_radius = 10
                    if role == "evader":
                        painter.setPen(QPen(color, 2))
                        painter.setBrush(QBrush(Qt.NoBrush))
                    else:
                        painter.setPen(QPen(color.darker(), 2))
                        painter.setBrush(QBrush(color.lighter(150)))
                    painter.drawEllipse(
                        px - robot_radius,
                        py - robot_radius,
                        robot_radius * 2,
                        robot_radius * 2,
                    )
                    painter.setPen(QPen(Qt.black))
                    painter.setFont(QFont("Arial", 9, QFont.Bold))
                    role_symbol = (
                        "E" if role == "evader" else "P" if role == "pursuer" else "?"
                    )
                    painter.drawText(px - 15, py - 15, f"{robot['id']}{role_symbol}")
                    # 目标
                    if (
                        "target" in robot
                        and robot["target"] != "未设置"
                        and "target_node" in robot
                        and robot["target_node"] != "未设置"
                    ):
                        try:
                            target_node_id = int(robot["target_node"])
                            if 0 <= target_node_id < node_num:
                                tx, ty = self.map_manager.node_coordinates[
                                    target_node_id
                                ]
                            else:
                                target_str = robot["target"].strip("()")
                                tx, ty = map(float, target_str.split(","))
                            tpx, tpy = self.world_to_pixel(tx, ty)
                            painter.setPen(QPen(color, 2, Qt.DashLine))
                            painter.drawLine(px, py, tpx, tpy)
                            painter.setPen(QPen(color.darker(), 2))
                            painter.setBrush(QBrush(Qt.NoBrush))
                            target_point_radius = 5
                            painter.drawEllipse(
                                tpx - target_point_radius,
                                tpy - target_point_radius,
                                target_point_radius * 2,
                                target_point_radius * 2,
                            )
                        except:
                            pass
                    # 状态指示器
                    status = robot.get("status", "unknown")
                    status_color = (
                        Qt.green
                        if status == "active"
                        else (
                            Qt.yellow
                            if status == "connecting"
                            else Qt.red if status == "disconnected" else Qt.gray
                        )
                    )
                    painter.setPen(QPen(status_color, 2))
                    painter.drawEllipse(
                        px - robot_radius - 3,
                        py - robot_radius - 3,
                        (robot_radius + 3) * 2,
                        (robot_radius + 3) * 2,
                    )
                # 标题
                painter.setPen(QPen(Qt.black))
                painter.setFont(QFont("Arial", 14, QFont.Bold))
                painter.drawText(20, 30, "追捕任务 - 实时监控")
                painter.setFont(QFont("Arial", 12))
                painter.drawText(20, 60, f"捕获次数: {self.capture_count}")

    def resizeEvent(self, event):
        """大小改变"""
        super().resizeEvent(event)
        if self.map_pixmap is not None:
            self.update_map_view()
        else:
            self.update_graph_view()

    def update_robots(self, robots, capture_count=0, belief_pr=None, belief_pos=None):
        """更新机器人和信念"""
        self.robots = robots
        self.belief_pr = belief_pr
        self.belief_pos = belief_pos
        self.capture_count = capture_count
        self.update()


if __name__ == "__main__":
    import sys
    from PyQt5.QtWidgets import QApplication, QMainWindow

    sys.path.append("..")
    from map_manager import MapManager

    # 创建应用程序
    app = QApplication(sys.argv)

    # 创建地图管理器并加载地图
    map_manager = MapManager()

    # 创建主窗口
    window = QMainWindow()
    window.setWindowTitle("地图显示测试")
    window.resize(1000, 800)

    # 创建地图控件并添加到窗口
    map_widget = GraphMapWidget(map_manager)
    window.setCentralWidget(map_widget)

    # 显示窗口
    window.show()

    # 运行应用程序
    sys.exit(app.exec_())
